from contextlib import contextmanager
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from models import Base, OrderRecord, PositionSnapshot, TradeSignal

class DatabaseManager:
    def __init__(self, db_url: str, echo=False):
        self.engine = create_engine(
            db_url,
            echo=echo,  # 显示SQL日志
            pool_size=20,
            max_overflow=10,
            pool_recycle=3600
        )
        self.Session = sessionmaker(bind=self.engine)
        Base.metadata.create_all(self.engine)  # 自动建表

    @contextmanager
    def session_scope(self):
        """提供事务范围的上下文管理器"""
        session = self.Session()
        try:
            yield session
            session.commit()
        except Exception as e:
            session.rollback()
            raise e
        finally:
            session.close()

    # ----- 订单操作 -----
    def save_order(self, order_data: dict):
        """保存订单记录"""
        with self.session_scope() as session:
            order = OrderRecord(
                order_id=order_data['order_id'],
                symbol=order_data['symbol'],
                side=order_data['side'],
                order_type=order_data.get('order_type'),
                quantity=order_data['quantity'],
                price=order_data.get('price'),
                status=order_data.get('status', 'SUBMITTED'),
                metadata=order_data.get('metadata')
            )
            session.add(order)

    def update_order_status(self, order_id: str, new_status: str):
        """更新订单状态"""
        with self.session_scope() as session:
            order = session.query(OrderRecord)\
                .filter(OrderRecord.order_id == order_id)\
                .first()
            if order:
                order.status = new_status
                order.updated_at = datetime.utcnow()

    # ----- 仓位快照 -----
    def save_position_snapshot(self, symbol: str, position_data: dict):
        """保存仓位快照"""
        with self.session_scope() as session:
            snapshot = PositionSnapshot(
                symbol=symbol,
                confirmed_long=position_data['confirmed_long'],
                confirmed_short=position_data['confirmed_short'],
                pending_long=position_data['pending_long'],
                pending_short=position_data['pending_short'],
                avg_price=position_data['avg_price']
            )
            session.add(snapshot)

    # ----- 交易信号 -----
    def log_trade_signal(self, signal_data: dict):
        """记录交易信号"""
        with self.session_scope() as session:
            signal = TradeSignal(
                signal_id=signal_data['signal_id'],
                symbol=signal_data['symbol'],
                signal_type=signal_data['signal_type'],
                entry_price=signal_data.get('entry_price'),
                stop_loss=signal_data.get('stop_loss'),
                volatility=signal_data.get('volatility'),
                executed=signal_data.get('executed', 'PENDING')
            )
            session.add(signal)

    # ----- 查询接口 -----
    def get_order_history(self, symbol: str, limit=100) -> list:
        """获取标的订单历史"""
        with self.session_scope() as session:
            return session.query(OrderRecord)\
                .filter(OrderRecord.symbol == symbol)\
                .order_by(OrderRecord.created_at.desc())\
                .limit(limit)\
                .all()

    def get_latest_position(self, symbol: str) -> PositionSnapshot:
        """获取最新仓位快照"""
        with self.session_scope() as session:
            return session.query(PositionSnapshot)\
                .filter(PositionSnapshot.symbol == symbol)\
                .order_by(PositionSnapshot.timestamp.desc())\
                .first()